{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4kPKrMTP_n_j"
   },
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "<h1 align=\"center\">Setting up Sessions</h1>\n",
    "\n",
    "A Session is a sequence of traces representing a single session (e.g. a session or a thread). Each response is represented as it's own trace, but these traces are linked together by being part of the same session.\n",
    "To associate traces together, you need to pass in a special metadata key where the value is the unique identifier for that thread.\n",
    "\n",
    "In this tutorial we will setup sessions using OpenAI and OpenInference instrumentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9y1DA5uu_n_k"
   },
   "source": [
    "## 1. Install Dependencies and Import Libraries\n",
    "\n",
    "Install dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"openai>=1.0.0\" arize-phoenix jsonschema openinference-instrumentation-openai openinference-instrumentation opentelemetry-api opentelemetry-sdk openinference-semantic-conventions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u9xTBlt__n_l"
   },
   "source": [
    "## Configure Your OpenAI API Key and Instantiate Your OpenAI Client\n",
    "\n",
    "Set your OpenAI API key if it is not already set as an environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "from openai import OpenAI\n",
    "\n",
    "if not (openai_api_key := os.getenv(\"OPENAI_API_KEY\")):\n",
    "    openai_api_key = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "client = OpenAI(api_key=openai_api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hY0NliPg_n_l"
   },
   "source": [
    "## Instrument Your OpenAI Client\n",
    "\n",
    "Instrument your OpenAI client with a tracer that emits telemetry data in OpenInference format. [OpenInference](https://arize-ai.github.io/open-inference-spec/trace/) is an open standard for capturing and storing LLM application traces that enables LLM applications to seamlessly integrate with LLM observability solutions such as Phoenix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openinference.instrumentation.openai import OpenAIInstrumentor\n",
    "\n",
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(project_name=\"openai-sessions-example\")\n",
    "OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MhVjxN_R_n_l"
   },
   "source": [
    "## Run Phoenix in the Background\n",
    "\n",
    "Launch Phoenix as a background session to collect the trace data emitted by your instrumented OpenAI client. Note that Phoenix should be run in a container in a production environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import phoenix as px\n",
    "\n",
    "px.launch_app().view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VoLsIY9o_n_m"
   },
   "source": [
    "## Create a bare-bones Agent\n",
    "\n",
    "Let's create the outline of an agent that will leverage OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "from openinference.instrumentation import using_session\n",
    "from openinference.semconv.trace import SpanAttributes\n",
    "from opentelemetry import trace\n",
    "\n",
    "session_id = str(uuid.uuid4())\n",
    "\n",
    "tracer = trace.get_tracer(__name__)\n",
    "\n",
    "\n",
    "@tracer.start_as_current_span(\n",
    "    name=\"agent\", attributes={SpanAttributes.OPENINFERENCE_SPAN_KIND: \"agent\"}\n",
    ")\n",
    "def assistant(\n",
    "    messages: list[dict],\n",
    "    session_id: str,\n",
    "):\n",
    "    current_span = trace.get_current_span()\n",
    "    current_span.set_attribute(SpanAttributes.SESSION_ID, session_id)\n",
    "    current_span.set_attribute(SpanAttributes.INPUT_VALUE, messages[-1].get(\"content\"))\n",
    "\n",
    "    # Propagate the session_id down to spans crated by the OpenAI instrumentation\n",
    "    # This is not strictly necessary, but it helps to correlate the spans to the same session\n",
    "    with using_session(session_id):\n",
    "        response = (\n",
    "            client.chat.completions.create(\n",
    "                model=\"gpt-3.5-turbo\",\n",
    "                messages=[{\"role\": \"system\", \"content\": \"You are a helpful assistant.\"}] + messages,\n",
    "            )\n",
    "            .choices[0]\n",
    "            .message\n",
    "        )\n",
    "\n",
    "    current_span.set_attribute(SpanAttributes.OUTPUT_VALUE, response.content)\n",
    "    return response\n",
    "\n",
    "\n",
    "messages = [{\"role\": \"user\", \"content\": \"hi! im bob\"}]\n",
    "response = assistant(\n",
    "    messages,\n",
    "    session_id=session_id,\n",
    ")\n",
    "messages = messages + [response, {\"role\": \"user\", \"content\": \"what's my name?\"}]\n",
    "response = assistant(\n",
    "    messages,\n",
    "    session_id=session_id,\n",
    ")\n",
    "messages = messages + [response, {\"role\": \"user\", \"content\": \"what's 4+5?\"}]\n",
    "response = assistant(\n",
    "    messages,\n",
    "    session_id=session_id,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
